#include <stdio.h>
#include <math.h>
#include <string.h>
#include "slave.h"

__thread_local volatile unsigned long get_reply,put_reply;
__thread_local double *A_lhs_array = NULL;//[BLKX][BLKY][BLKZ];
__thread_local double *A_rhs_array = NULL;//[BLKX+2][BLKY+2][BLKZ+2];

extern double A[2][SIZEX][SIZEY][SIZEZ];
extern int thread_num, iter, N, ntilex, ntiley, ntilez;
extern int startk[64], startj[64], starti[64];

#define alpha_d 0.0876
#define beta_d  0.0765
#define A_lhs(i,j,k) A_lhs_array[(i)*(cur_blky)*(cur_blkz)+(j)*(cur_blkz)+(k)]
#define A_rhs(i,j,k) A_lhs_array[(i)*(cur_blky+2)*(cur_blkz+2)+(j)*(cur_blkz+2)+(k)]



void func()
{
   int i,j,k,n,ii,jj,kk,my_id;//N ~ average num tiles that each thread assigned
   int dest,src,my_tiles;
   int cur_blkx, cur_blky, cur_blkz;
   my_id = athread_get_id(-1);
   //printf("my_id = %d \n",my_id);
   dest = iter % 2 ;
   src = 1-dest;
   n = 0;
   if(my_id < thread_num-1)
      my_tiles = N;
   else
      my_tiles = ntilex*ntiley*ntilez-N*(thread_num-1);
   for(i = starti[my_id]; n < my_tiles && i < (SIZEX-2)-BLKX+1; i += BLKX)
     for(j = startj[my_id]; n < my_tiles && j < (SIZEY-2)-BLKY+1; j += BLKY)
       for(k = startk[my_id]; n < my_tiles && k < (SIZEZ-2)-BLKZ+1 ; k += BLKZ)
       {
         int stride,tsize,bsize;
         int dmai, lhs_size, rhs_size;
         ++n;
         //average tile size: <SIZEX, BLXY,BLKZ>
         cur_blkx = BLKX;
         if(i+BLKX+1 > SIZEX) //last tile at dim X
            cur_blkx = ((SIZEX-2)%BLKX);
         cur_blky = BLKY;
         if(j+BLKY+1 > SIZEY) //last tile at dim Y
            cur_blky = ((SIZEY-2)%BLKY);
         cur_blkz = BLKZ;
         if(k+BLKZ+1 > SIZEZ) //last tile at dim Z
            cur_blkz = ((SIZEZ-2)%BLKZ);
         lhs_size = cur_blkx*cur_blky*cur_blkz;
         rhs_size = (cur_blkx+2)*(cur_blky+2)*(cur_blkz+2);
         if(A_lhs_array == NULL || A_rhs_array == NULL )
         {   
           A_lhs_array = (double *)ldm_malloc(lhs_size * sizeof(double));
           A_rhs_array = (double *)ldm_malloc(rhs_size * sizeof(double));
         }
         //update data
         get_reply = 0;
         for(dmai = i; dmai < i+cur_blkx+2; dmai++)
         {
           stride = (SIZEZ-(cur_blkz+2))*8;
           tsize = (cur_blky+2)*(cur_blkz+2)*8;
           bsize = (cur_blkz+2)*8;
           athread_get(PE_MODE,&A[src][dmai][j][k]
                 , &A_rhs(dmai-i,0,0)
                 , tsize, &get_reply
                 , 0, stride, bsize);
         }
         while(get_reply!=cur_blkx+2); 
        /* ii = 1;
         jj = 1;
         kk = 1;
         printf("%d\n",k);
         printf("slave:%20.20f,%20.20f,%20.20f,%20.20f,%20.20f,%20.20f,%20.20f\n",A_rhs[ii][jj][kk],A_rhs[ii-1][jj][kk],A_rhs[ii+1][jj][kk],A_rhs[ii][jj-1][kk],A_rhs[ii][jj+1][kk],A_rhs[ii][jj][kk-1],A_rhs[ii][jj][kk+1]);*/
         //compute 
         for(ii = 1; ii < cur_blkx+1; ++ii) 
           for(jj = 1; jj < cur_blky+1; ++jj) 
            for(kk = 1; kk < cur_blkz+1; ++kk)  
                A_lhs((ii-1),(jj-1),(kk-1)) = alpha_d * (A_rhs(ii,jj,kk)) + \
                        beta_d * (A_rhs(ii-1,jj,kk) + A_rhs(ii+1,jj,kk) +\
                        A_rhs(ii,jj-1,kk) + A_rhs(ii,jj+1,kk) +\
                        A_rhs(ii,jj,kk-1) + A_rhs(ii,jj,kk+1));
         //printf("slave:%20.20f\n",A_lhs[0][0][0]);
         //write back
         put_reply=0;
         for(dmai = i; dmai < i+cur_blkx; dmai++)
         {
           stride = (SIZEZ-cur_blkz)*8;
           bsize = cur_blkz*8;
           tsize = bsize*cur_blky;
           athread_put(PE_MODE,&A_lhs(dmai-i,0,0)
                     , &A[dest][dmai+1][j+1][k+1]
                     ,tsize,&put_reply,stride,bsize);
         }
         while(put_reply!=cur_blkx);
         ldm_free(A_lhs_array,lhs_size*sizeof(double));
         ldm_free(A_rhs_array,rhs_size*sizeof(double));
       }
}
